package commons.cs.cop;

import commons.util.ExceptionUtil;
import commons.util.base.Equals;
import org.jgrapht.DirectedGraph;
import org.jgrapht.EdgeFactory;
import org.jgrapht.alg.EdmondsKarpMaximumFlow;
import org.jgrapht.graph.SimpleDirectedGraph;

import java.util.*;

/**
 * Consider a problem when you have more workers and less jobs. You have a mapping of workers and
 * jobs (depending on who can do what). You want to maximize the number of jobs assigned with the
 * condition that a worker is not assigned more than one job. Note that a job might have multiple
 * workers.
 *
 * @version 1.2.3
 * @author Susanta Tewari
 * @history Created on 3/4/13
 * @since 1.2.3
 * @param <J>
 * @param <W>
 */
public class JobAssignment<J, W> {

    private Equals<W> equalsImpl;

    /**
     * The semantics of worker equality can be extended with an implementation of {@code Equals}.
     *
     * @param equalsImpl an equals implementation for the workers
     */
    public void setEqualsImpl(Equals equalsImpl) {
        this.equalsImpl = equalsImpl;
    }

    /**
     * Uses maximum bipartite matching to solve this problem. Max flow algo (Edmonds-Karp) will
     * be used.
     *
     * @return the desired assignment
     * @throws NullPointerException if {@code input} is {@code null}
     * @throws IllegalArgumentException if {@code input} is empty
     */
    public Map<J, Set<W>> assign(Map<J, Set<W>> input) {

        Objects.requireNonNull(input);
        ExceptionUtil.throwArgEx(input.isEmpty(), "input is empty");

        final String src                                = "s",
                     sink                               = "t";
        final DirectedGraph<Object, MyEdge> graph       = createNetwork(src, sink, input);
        EdmondsKarpMaximumFlow<Object, MyEdge> flowAlgo = new EdmondsKarpMaximumFlow<>(graph);

        flowAlgo.calculateMaximumFlow(src, sink);

        final Map<MyEdge, Double> maxFlow = flowAlgo.getMaximumFlow();
        Map<J, Set<W>> output             = new HashMap<>();

        for (MyEdge e : maxFlow.keySet()) {

            if (e.src == src) continue;
            if (e.target == sink) continue;

            final Object s = e.src;

            if (!output.containsKey(s)) output.put((J) s, new HashSet<W>());

            boolean connected = maxFlow.get(e) == 1.0;

            if (connected) output.get(s).add((W) e.target);
        }

        return output;
    }

    private DirectedGraph<Object, MyEdge> createNetwork(String src, String sink,
            Map<J, Set<W>> input) {

        final DirectedGraph<Object, MyEdge> graph =
            new SimpleDirectedGraph<>(new EdgeFactory<Object, MyEdge>() {

            @Override
            public MyEdge createEdge(final Object sourceVertex, final Object targetVertex) {
                return new MyEdge(sourceVertex, targetVertex);
            }

        });

        graph.addVertex(src);
        graph.addVertex(sink);

        for (J j : input.keySet()) {

            graph.addVertex(j);
            graph.addEdge(src, j);

            for (W w : input.get(j)) {

                final boolean added = graph.addVertex(w);

                if (added) graph.addEdge(w, sink);

                graph.addEdge(j, w);
            }
        }

        return graph;
    }

    private static class MyEdge {

        private Object src, target;

        private MyEdge(Object src, Object target) {

            this.src    = src;
            this.target = target;
        }

        @Override
        public String toString() {
            return src + "->" + target;
        }
    }
}
